import gym
import numpy as np
import torch
import os

def create_env(env_name, seed):
    env = gym.make(env_name)
    env.seed(seed)
    env.action_space.np_random.seed(seed)
    return env

def init_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed_all(seed)

def make_dir(*path_parts):
    dir_path = os.path.join(*path_parts)
    try:
        os.mkdir(dir_path)
    except OSError:
        pass
    return dir_path

def save_model(agent, args):
    make_dir("./save_model")
    model_name = 'algname_' + args.alg + '-' + 'envname_' + args.env + '-' + 'seed_' + str(args.seed) + '-'   \
                    + 'Qnmus_' + str(args.Qnums) + '-' + 'minqnum_' + str(args.M) + '-' + 'UTD_' + str(args.UTD) + '-'   \
                    + 'policynum_' + str(args.policynum) + '-' + str(args.lambd) + 'batchsize_' + str(args.batchsize) + \
                    '-' + args.label + '-actor' + '.pkl'
    torch.save(agent.actor.state_dict(), './save_model/' + model_name)
    model_name = 'algname_' + args.alg + '-' + 'envname_' + args.env + '-' + 'seed_' + str(args.seed) + '-'   \
                    + 'Qnmus_' + str(args.Qnums) + '-' + 'minqnum_' + str(args.M) + '-' + 'UTD_' + str(args.UTD) + '-'   \
                    + 'policynum_' + str(args.policynum) + '-' + str(args.lambd) + 'batchsize_' + str(args.batchsize) + \
                    '-' + args.label + '-critic' + '.pkl'
    for q_i in range(agent.Qnums):
        torch.save(agent.Q_network_list[q_i].state_dict(), './save_model/' + str(q_i) + model_name)
    for q_i in range(agent.Qnums):
        torch.save(agent.Q_target_network_list[q_i].state_dict(), './save_model/' + "target" + str(q_i) + model_name)